from langchain.tools import BaseTool

from langchain.chains import RetrievalQA
from langchain_neo4j import Neo4jVector
from pydantic import BaseModel, Field, ConfigDict
from typing import Any, Type, Optional
from dotenv import load_dotenv
from model.my_chat_model import ChatModel
import os


# 定义输入参数的数据模型类
class DocumentInput(BaseModel):
    question: str = Field(..., description="问题")


# 定义工具类
class DocumentTool(BaseTool):
    # 定义模型，是否允许输入参数，输入参数的数据模型类
    model_config = ConfigDict(arbitrary_types_allowed=True)

    def __init__(self, **kwargs: Any):
        super().__init__(
            name="get_document_tool",
            description="主要用于查询文档类型数据，例如(动物解释，动物说明)，必须输入的参数是问题",
            **kwargs
        )

    # 定义工具参数
    args_schema: Type[BaseModel] = DocumentInput

    # 定义工具方法
    def _run(self, question: str):
        load_dotenv()
        chat = ChatModel()
        embedding_model = chat.get_embedding_model()
        # 查询存入到neo4j的文档
        vector_sore = Neo4jVector(
            embedding=embedding_model,
            url=os.getenv("NEO4J_URL"),
            username=os.getenv("NEO4J_USERNAME"),
            password=os.getenv("NEO4J_PASSWORD"),
            index_name="企业",  # 索引名称
            node_label="企业",  # 节点标签
            text_node_property="text",  # 节点属性
            embedding_node_property="embedding"
        )
        # 获取聊天大模型
        llm = chat.get_line_model()
        # 创建链
        chain = RetrievalQA.from_chain_type(
            llm=llm,
            retriever=vector_sore.as_retriever(searc_kwargs={"k": 3}),  #使用向量搜索, 搜索的向量索引名称
            verbose=True,  # 是否打印日志 (可选)
            return_source_documents=True,  # 是否返回源文档(可选)
        )
        rs = chain.invoke(question)
        return rs["result"]
